# Symbol 是解方程的库
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sympy import Symbol, Eq, solve, summation
from mpl_toolkits.mplot3d import Axes3D


plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def solve_eq_1_10(v_set):
    # 1
    eq_list0 = []
    steady_state_rate = Symbol('steady_state_rate')

    number_of_capital_firm_workers = Symbol('number_of_capital_firm_workers')
    output_of_capital = Symbol('output_of_capital')
    labor_productivity = Symbol('labor_productivity')
    eq_list0.append(Eq(number_of_capital_firm_workers, output_of_capital / labor_productivity))
    # 4个未知数，1个方程

    unit_cost_of_capital_firm = Symbol('unit_cost_of_capital_firm')
    households_wage = Symbol('households_wage')
    eq_list0.append(Eq(unit_cost_of_capital_firm, households_wage / labor_productivity))
    # 6个未知数 2个方程

    # eq3
    price_of_capital = Symbol('price_of_capital')
    mark_up_of_capital = Symbol('mark_up_of_capital')
    eq_list0.append(Eq(price_of_capital, (1 + mark_up_of_capital) * unit_cost_of_capital_firm))
    # 8个未知数 3个方程

    # eq4
    deposits_capital_firms_holds = Symbol('deposits_capital_firms_holds')
    ratio_of_deposits_to_wage = Symbol('ratio_of_deposits_to_wage')
    eq_list0.append(
        Eq(deposits_capital_firms_holds, ratio_of_deposits_to_wage * households_wage * number_of_capital_firm_workers)
    )
    # 10个未知数 4 个方程

    # eq5
    capital_of_consumption_firm = Symbol('capital_of_consumption_firm')
    duration_or_capital = Symbol('duration_or_capital')
    eq_list0.append(Eq(output_of_capital, capital_of_consumption_firm / duration_or_capital))
    # 12个未知数 5个方程

    # eq6
    capital_firm_inventory = Symbol('capital_firm_inventory')
    ratio_of_inventory_to_output = Symbol('ratio_of_inventory_to_output')
    eq_list0.append(Eq(capital_firm_inventory, ratio_of_inventory_to_output * output_of_capital))
    # 14个未知数 6个方程

    # eq7
    profits_of_capital_firm_before_tax = Symbol('profits_of_capital_firm_before_tax')
    profits_of_capital_firm_after_tax = Symbol('profits_of_capital_firm_after_tax')

    interest_rate_of_deposits = Symbol('interest_rate_of_deposits')
    interest_rate_of_loans = Symbol('interest_rate_of_loans')
    loans_of_capital_firms = Symbol('loans_of_capital_firms')
    deposits_interests_of_capital_firm = Symbol('deposits_interests_of_capital_firm')
    loan_interests_of_capital_firm = Symbol('loan_interests_of_capital_firm')
    change_inventory_of_capital_firm = Symbol('change_inventory_of_capital_firm')

    oo1_steady_state_rate = Symbol('oo1_steady_state_rate')
    oo2_steady_state_rate = Symbol('oo2_steady_state_rate')
    tax_of_capital_firms = Symbol('tax_of_capital_firms')
    tax_rate = Symbol('tax_rate')

    eq_list0.append(Eq(oo1_steady_state_rate, 1 / (steady_state_rate + 1)))
    eq_list0.append(Eq(oo2_steady_state_rate, steady_state_rate / (steady_state_rate + 1)))
    eq_list0.append(
        Eq(deposits_interests_of_capital_firm,
           interest_rate_of_deposits * deposits_capital_firms_holds * oo1_steady_state_rate)
    )
    eq_list0.append(
        Eq(loan_interests_of_capital_firm, loans_of_capital_firms * interest_rate_of_loans * oo1_steady_state_rate))
    eq_list0.append(Eq(change_inventory_of_capital_firm,
                       capital_firm_inventory * unit_cost_of_capital_firm * oo2_steady_state_rate))
    eq_list0.append(Eq(profits_of_capital_firm_before_tax,
                       price_of_capital * output_of_capital +
                       deposits_interests_of_capital_firm +
                       change_inventory_of_capital_firm -
                       households_wage * number_of_capital_firm_workers -
                       loan_interests_of_capital_firm
                       ))
    eq_list0.append(Eq(tax_of_capital_firms, tax_rate * profits_of_capital_firm_before_tax))
    eq_list0.append(Eq(profits_of_capital_firm_after_tax, profits_of_capital_firm_before_tax - tax_of_capital_firms))
    # 26个未知数，14个方程

    # eq9
    dividends_of_capital_firms = Symbol('dividends_of_capital_firms')
    ratio_of_dividends_to_profits_left = Symbol('ratio_of_dividends_to_profits_left')
    eq_list0.append(Eq(dividends_of_capital_firms,
                       ratio_of_dividends_to_profits_left * (profits_of_capital_firm_before_tax - tax_of_capital_firms)
                       ))
    # 28个未知数，16个方程

    # eq10
    nominal_output_of_capital = Symbol('nominal_output_of_capital')
    delta_loans_of_capital_firms = Symbol('delta_loans_of_capital_firms')
    eq_list0.append(Eq(delta_loans_of_capital_firms, loans_of_capital_firms * oo2_steady_state_rate))
    eq_list0.append(Eq(nominal_output_of_capital, output_of_capital * price_of_capital))
    eq_list0.append(Eq(delta_loans_of_capital_firms,
                       capital_firm_inventory * unit_cost_of_capital_firm * oo2_steady_state_rate +
                       deposits_capital_firms_holds * oo2_steady_state_rate -
                       (profits_of_capital_firm_before_tax - tax_of_capital_firms - dividends_of_capital_firms)
                       ))
    # 30个未知数，19个方程

    variable_collector0 = {
        'steady_state_rate': (steady_state_rate, "稳态增速", "外生设定"),  # 1
        'number_of_capital_firm_workers': (number_of_capital_firm_workers, "资本公司员工数量", "外生设定"),  # 2
        'output_of_capital': (output_of_capital, "资本品的产出", "-"),  # 3
        'labor_productivity': (labor_productivity, "生产资本品劳动生产率", "-"),  # 4
        'unit_cost_of_capital_firm': (unit_cost_of_capital_firm, "资本品单位生产成本", "-"),  # 5
        'households_wage': (households_wage, "员工工资", "外生设定"),  # 6
        'price_of_capital': (price_of_capital, "资本品价格", "-"),  # 7
        'mark_up_of_capital': (mark_up_of_capital, "资本加价比例", "外生设定"),  # 8
        'deposits_capital_firms_holds': (deposits_capital_firms_holds, "资本品公司存款", "-"),  # 9
        'ratio_of_deposits_to_wage': (ratio_of_deposits_to_wage, "公司持有存款占应付工资比例", "外生设定"),  # 10
        'capital_of_consumption_firm': (capital_of_consumption_firm, "消费品公司所持有的资本", "-"),  # 11
        'duration_or_capital': (duration_or_capital, "资本折旧年限", "外生设定"),  # 12
        'capital_firm_inventory': (capital_firm_inventory, "资本品公司存货", "-"),  # 13
        'ratio_of_inventory_to_output': (ratio_of_inventory_to_output, "资本品公司存货比例", "外生设定"),  # 14
        'profits_of_capital_firm_before_tax': (profits_of_capital_firm_before_tax, "资本品公司税前利润", "-"),  # 15
        'interest_rate_of_deposits': (interest_rate_of_deposits, "存款利率", "外生设定"),  # 16
        'interest_rate_of_loans': (interest_rate_of_loans, "贷款利率", "外生设定"),  # 17
        'loans_of_capital_firms': (loans_of_capital_firms, "资本品公司贷款", "-"),  # 18
        'deposits_interests_of_capital_firm': (deposits_interests_of_capital_firm, "资本品公司存款利息", "-"),  # 19
        'loan_interests_of_capital_firm': (loan_interests_of_capital_firm, "资本品公司贷款利息", "-"),  # 20
        'change_inventory_of_capital_firm': (change_inventory_of_capital_firm, "资本品公司存货增加", "-"),  # 21
        'tax_of_capital_firms': (tax_of_capital_firms, "资本品公司税收", "-"),
        'profits_of_capital_firm_after_tax': (profits_of_capital_firm_after_tax, "资本品公司税后利润", "-"),
        'tax_rate': (tax_rate, "税率", "外生设定"),
        'dividends_of_capital_firms': (dividends_of_capital_firms, "资本品公司股利", "-"),
        'ratio_of_dividends_to_profits_left': (ratio_of_dividends_to_profits_left, "公司股利比率", "外生设定"),
        'delta_loans_of_capital_firms': (delta_loans_of_capital_firms, "资本品公司贷款增加", "-"),
        'oo1_steady_state_rate': (oo1_steady_state_rate, "辅助参数1", "外生设定"),
        'oo2_steady_state_rate': (oo2_steady_state_rate, "辅助参数2", "外生设定"),
        'nominal_output_of_capital': (nominal_output_of_capital, "资本品名义产出", '-'),
    }

    eq_list0.append(Eq(steady_state_rate, v_set['steady_state_rate']))  # 平稳状态下增长率
    eq_list0.append(Eq(number_of_capital_firm_workers, 1000))  # 资本品生产部门工人数量
    eq_list0.append(Eq(households_wage, v_set['households_wage']))  # 工人工资
    eq_list0.append(Eq(mark_up_of_capital, 0.075))  # 资本品加价比例
    eq_list0.append(Eq(ratio_of_deposits_to_wage, 1))
    eq_list0.append(Eq(capital_of_consumption_firm, 40000))
    eq_list0.append(Eq(duration_or_capital, 20))
    eq_list0.append(Eq(ratio_of_inventory_to_output, 0.1))
    eq_list0.append(Eq(interest_rate_of_deposits, 0.0025))
    eq_list0.append(Eq(interest_rate_of_loans, 0.0075))
    eq_list0.append(Eq(tax_rate, v_set['tax_rate']))
    eq_list0.append(Eq(ratio_of_dividends_to_profits_left, 0.9))

    variables = [v[0] for k, v in variable_collector0.items()]
    print(f"一共有未知数{len(variables)}个, 方程{len(eq_list0)}个")
    results = solve(eq_list0, variables)
    solved_variables = dict(zip([k for k, v in variable_collector0.items()], results[0]))
    set_variables = {k: v for k, v in solved_variables.items() if variable_collector0[k][2] == '外生设定'}
    return solved_variables, set_variables, variable_collector0


def solve_eq_11_22(solved_variables, variable_collector0):
    eq_list = []
    number_of_consumption_firm_workers = Symbol('number_of_consumption_firm_workers')
    output_of_consumption = Symbol('output_of_consumption')
    capital_productivity = Symbol('capital_productivity')
    labor_capital_ratio = Symbol('labor_capital_ratio')
    eq_list.append(
        Eq(number_of_consumption_firm_workers, output_of_consumption / (capital_productivity * labor_capital_ratio))
    )
    # 4个未知数，1个方程

    unit_variable_cost_of_consumption_firm = Symbol('unit_variable_cost_of_consumption_firm')
    households_wage = variable_collector0['households_wage'][0]
    eq_list.append(
        Eq(unit_variable_cost_of_consumption_firm,
           households_wage * number_of_consumption_firm_workers / output_of_consumption)
    )
    # 5个未知数， 2个方程
    unit_cost_of_consumption_firm = Symbol('unit_cost_of_consumption_firm')
    nominal_output_of_consumpiton = Symbol('nominal_output_of_consumption')
    supp_para_for_consumption_unit_cost = Symbol("supp_para_for_consumption_unit_cost")

    capital_of_consumption_firm = variable_collector0['capital_of_consumption_firm'][0]
    duration_or_capital = variable_collector0['duration_or_capital'][0]
    price_of_capital = variable_collector0['price_of_capital'][0]
    eq_list.append(
        Eq(unit_cost_of_consumption_firm,
           (
                   households_wage * number_of_consumption_firm_workers + supp_para_for_consumption_unit_cost * price_of_capital *
                   capital_of_consumption_firm / (duration_or_capital ** 2)) / output_of_consumption)
    )
    # 8个未知数，3个方程

    supp_para_for_stock_of_capital = Symbol("supp_para_for_stock_of_capital")
    stock_of_capital_in_consumption_firm = Symbol("stock_of_capital_in_consumption_firm")
    price_of_capital = variable_collector0['price_of_capital'][0]
    eq_list.append(
        Eq(stock_of_capital_in_consumption_firm,
           supp_para_for_stock_of_capital * price_of_capital * capital_of_consumption_firm / (duration_or_capital ** 2))
    )
    # 10个未知数，4个方程

    mark_up_of_consumption = Symbol("mark_up_of_consumption")
    price_of_consumption = Symbol("price_of_consumption")

    eq_list.append(
        Eq(price_of_consumption, (1 + mark_up_of_consumption) * unit_variable_cost_of_consumption_firm)
    )
    eq_list.append(
        Eq(nominal_output_of_consumpiton, output_of_consumption * price_of_consumption)
    )
    # 12 个未知数 6个方程

    deposits_consumption_firm_holds = Symbol('deposits_consumption_firm_holds')
    ratio_of_deposits_to_wage = variable_collector0['ratio_of_deposits_to_wage'][0]
    eq_list.append(
        Eq(deposits_consumption_firm_holds,
           ratio_of_deposits_to_wage * households_wage * number_of_consumption_firm_workers)
    )
    # 13个未知数 7个方程

    capacity_of_utilization_of_capital = Symbol("capacity_of_utilization_of_capital")
    eq_list.append(
        Eq(output_of_consumption,
           capital_of_consumption_firm * capacity_of_utilization_of_capital * capital_productivity)
    )
    # 14个未知数, 8个方程
    ratio_of_inventory_to_output = variable_collector0['ratio_of_inventory_to_output'][0]
    consumption_firm_inventory = Symbol("consumption_firm_inventory")
    eq_list.append(
        Eq(consumption_firm_inventory, ratio_of_inventory_to_output * output_of_consumption)
    )
    # 15个未知数，9个方程

    profits_of_consumption_firm = Symbol("profits_of_consumption_firm")
    loans_of_consumption_firms = Symbol("loans_of_consumption_firms")

    deposits_interests_of_consumption_firm = Symbol('deposits_interests_of_consumption_firm')

    amortization_of_consumption_firms = Symbol("amortization_of_consumption_firms")
    eq_list.append(
        Eq(amortization_of_consumption_firms,
           supp_para_for_consumption_unit_cost * price_of_capital * capital_of_consumption_firm / (
                   duration_or_capital ** 2)
           ))
    interest_rate_of_deposits = variable_collector0['interest_rate_of_deposits'][0]
    oo1_steady_state_rate = variable_collector0['oo1_steady_state_rate'][0]
    oo2_steady_state_rate = variable_collector0['oo2_steady_state_rate'][0]
    eq_list.append(
        Eq(deposits_interests_of_consumption_firm,
           interest_rate_of_deposits * deposits_consumption_firm_holds * oo1_steady_state_rate)
    )

    interest_rate_of_loans = variable_collector0['interest_rate_of_loans'][0]

    eq_list.append(
        Eq(profits_of_consumption_firm,
           price_of_consumption * output_of_consumption +
           deposits_interests_of_consumption_firm +
           consumption_firm_inventory * unit_cost_of_consumption_firm * oo2_steady_state_rate -
           households_wage * number_of_consumption_firm_workers -
           interest_rate_of_loans * loans_of_consumption_firms * oo1_steady_state_rate -
           amortization_of_consumption_firms
           )
    )
    tax_rate = variable_collector0['tax_rate'][0]
    tax_of_consumption_firms = Symbol('tax_of_consumption_firms')
    eq_list.append(
        Eq(tax_of_consumption_firms, tax_rate * profits_of_consumption_firm)
    )
    # 20个未知数， 13个方程
    ratio_of_dividends_to_profits_left = variable_collector0['ratio_of_dividends_to_profits_left'][0]
    dividends_of_consumption_firms = Symbol('dividends_of_consumption_firms')
    eq_list.append(
        Eq(dividends_of_consumption_firms,
           ratio_of_dividends_to_profits_left * (profits_of_consumption_firm * (1 - tax_rate)
                                                 ))
    )
    output_of_capital = variable_collector0['output_of_capital'][0]
    eq_list.append(
        Eq(loans_of_consumption_firms * oo2_steady_state_rate,
           output_of_capital * price_of_capital +
           consumption_firm_inventory * unit_cost_of_consumption_firm * oo2_steady_state_rate +
           deposits_consumption_firm_holds * oo2_steady_state_rate -
           (profits_of_consumption_firm - tax_of_consumption_firms - dividends_of_consumption_firms) -
           amortization_of_consumption_firms
           )
    )

    variable_collector = {
        'number_of_consumption_firm_workers': (number_of_consumption_firm_workers, "消费品公司员工数量", "外生设定"),
        'output_of_consumption': (output_of_consumption, "消费品公司实际产出", "-"),
        'capital_productivity': (capital_productivity, "生产消费品的资本生产率", "外生设定"),
        'labor_capital_ratio': (labor_capital_ratio, "生产消费品的劳动资本比率", "外生设定"),
        'unit_variable_cost_of_consumption_firm': (unit_variable_cost_of_consumption_firm, "消费品单位可变生产成本", "-"),
        'unit_cost_of_consumption_firm': (unit_cost_of_consumption_firm, "消费品单位生产成本(包括折旧)", "-"),
        'supp_para_for_consumption_unit_cost': (supp_para_for_consumption_unit_cost, "用于计算消费品单位生产成本的系数", "-"),
        'supp_para_for_stock_of_capital': (supp_para_for_stock_of_capital, "用于计算资本存量的系数", "-"),
        'stock_of_capital_in_consumption_firm': (stock_of_capital_in_consumption_firm, "消费品公司资本存量", "-"),
        'mark_up_of_consumption': (mark_up_of_consumption, "资本公司加价比例", "外生设定"),
        'price_of_consumption': (price_of_consumption, "消费品价格", "-"),
        'nominal_output_of_consumpiton': (nominal_output_of_consumpiton, "消费品公司名义产出", "-"),
        'deposits_consumption_firm_holds': (deposits_consumption_firm_holds, "消费品公司存款", "-"),
        'deposits_interests_of_consumption_firm': (deposits_interests_of_consumption_firm, "消费品公司存款利息", "-"),
        'amortization_of_consumption_firms': (amortization_of_consumption_firms, "消费品公司折旧", "-"),
        'capacity_of_utilization_of_capital': (capacity_of_utilization_of_capital, "消费品公司资本利用率", "外生设定"),
        'consumption_firm_inventory': (consumption_firm_inventory, "消费品公司存货", "-"),
        'profits_of_consumption_firm': (profits_of_consumption_firm, "消费品公司利润", "-"),
        'loans_of_consumption_firms': (loans_of_consumption_firms, "消费品公司贷款", "-"),
        'tax_of_consumption_firms': (tax_of_consumption_firms, "消费品公司交税", "-"),
        'dividends_of_consumption_firms': (dividends_of_consumption_firms, "消费品公司股利分配", "-"),
    }

    steady_state_rate = variable_collector0['steady_state_rate'][0]

    eq_list.append(Eq(number_of_consumption_firm_workers, 5000))
    eq_list.append(Eq(capital_productivity, 1))
    # eq_list.append(Eq(labor_capital_ratio, 6.4))
    eq_list.append(
        Eq(supp_para_for_consumption_unit_cost, sum([(1.0 / (1.0 + steady_state_rate)) ** i for i in range(1, 21)])))
    eq_list.append(
        Eq(supp_para_for_stock_of_capital, sum([i / (1 + steady_state_rate) ** (20 - i) for i in range(1, 21)])))
    eq_list.append(Eq(mark_up_of_consumption, 0.319))
    eq_list.append(Eq(capacity_of_utilization_of_capital, 0.8))
    # 21个未知数（6个参数）， 15个方程
    for k, v in solved_variables.items():
        eq_list.append(Eq(variable_collector0[k][0], float(v)))
    variable_collector1 = {**variable_collector, **variable_collector0}
    variables = [v[0] for k, v in variable_collector1.items()]
    print(f"一共有未知数{len(variables)}个, 方程{len(eq_list)}个")
    results = solve(eq_list, variables)
    solved_variables = dict(zip([k for k, v in variable_collector1.items()], results[0]))
    solved_variables = {k: v for k, v in solved_variables.items()}
    solved_variables_output = {k: v for k, v in solved_variables.items() if k in variable_collector1}
    set_variables = {k: v for k, v in solved_variables.items() if variable_collector1[k][2] == '外生设定'}
    return solved_variables_output, set_variables, variable_collector1


def solve_eq_23_30(solved_variables, collector):
    # eq 23
    number_of_government_workers = Symbol('number_of_government_workers')
    number_of_total_employment_workers = Symbol('number_of_total_employment_workers')
    eq_list2 = []
    number_of_capital_firm_workers = collector['number_of_capital_firm_workers'][0]
    number_of_consumption_firm_workers = collector['number_of_consumption_firm_workers'][0]
    eq_list2.append(
        Eq(number_of_total_employment_workers,
           number_of_capital_firm_workers + number_of_consumption_firm_workers + number_of_government_workers)
    )
    # 变量2 ，公式1
    number_of_households = Symbol('number_of_households')
    ratio_of_dole_to_wage = Symbol('ratio_of_dole_to_wage')
    dividends_of_banks = Symbol('dividends_of_banks')
    net_income_of_households = Symbol('net_income_of_households')

    tax_of_households = Symbol('tax_of_households')
    wages_of_households = Symbol('wages_of_households')
    deposits_of_households = Symbol('deposits_of_households')
    deposits_interests_of_households = Symbol('deposits_interests_of_households')
    dividends_of_households = Symbol('dividends_of_households')
    households_income_before_tax_and_dole = Symbol('households_income_before_tax_and_dole')
    doles_of_households = Symbol('doles_of_households')

    households_wage = collector['households_wage'][0]
    interest_rate_of_deposits = collector['interest_rate_of_deposits'][0]
    oo1_steady_state_rate = collector['oo1_steady_state_rate'][0]
    dividends_of_capital_firms = collector['dividends_of_capital_firms'][0]
    dividends_of_consumption_firms = collector['dividends_of_consumption_firms'][0]
    tax_rate = collector['tax_rate'][0]

    eq_list2.append(
        Eq(wages_of_households, households_wage * number_of_total_employment_workers)
    )
    eq_list2.append(
        Eq(deposits_interests_of_households, deposits_of_households * interest_rate_of_deposits * oo1_steady_state_rate)
    )
    eq_list2.append(
        Eq(dividends_of_households, dividends_of_capital_firms + dividends_of_consumption_firms + dividends_of_banks)
    )
    eq_list2.append(
        Eq(households_income_before_tax_and_dole,
           wages_of_households + deposits_interests_of_households + dividends_of_households * oo1_steady_state_rate)
    )
    # eq 25
    eq_list2.append(
        Eq(tax_of_households, households_income_before_tax_and_dole * tax_rate)
    )
    eq_list2.append(
        Eq(doles_of_households,
           ratio_of_dole_to_wage * households_wage * (number_of_households - number_of_total_employment_workers))
    )
    # eq 24
    eq_list2.append(
        Eq(net_income_of_households,
           households_income_before_tax_and_dole - tax_of_households + doles_of_households
           )
    )
    # 变量13，公式8

    real_households_consumption = Symbol('real_households_consumption')
    net_wealth_of_households = Symbol('net_wealth_of_households')
    households_consumption_para_1 = Symbol('households_consumption_para_1')
    households_consumption_para_2 = Symbol('households_consumption_para_2')

    price_of_consumption = collector['price_of_consumption'][0]
    output_of_consumption = collector['output_of_consumption'][0]

    eq_list2.append(
        Eq(real_households_consumption,
           households_consumption_para_1 * net_income_of_households / price_of_consumption +
           households_consumption_para_2 * oo1_steady_state_rate * net_wealth_of_households / price_of_consumption
           )
    )
    # 变量17，公式9

    eq_list2.append(Eq(real_households_consumption, output_of_consumption))
    nominal_households_consumption = Symbol('nominal_households_consumption')
    eq_list2.append(Eq(nominal_households_consumption, real_households_consumption * price_of_consumption))
    # 变量18，公式11

    oo2_steady_state_rate = collector['oo2_steady_state_rate'][0]

    eq_list2.append(
        Eq(net_wealth_of_households * oo2_steady_state_rate,
           net_income_of_households - nominal_households_consumption - dividends_of_households * oo1_steady_state_rate + dividends_of_households
           )
    )
    eq_list2.append(Eq(net_wealth_of_households, deposits_of_households))
    # 变量18 ，公式13

    eq_list2.append(Eq(number_of_government_workers, 1360))
    eq_list2.append(Eq(number_of_households, 8000))
    eq_list2.append(Eq(ratio_of_dole_to_wage, 0.4))
    eq_list2.append(Eq(households_consumption_para_2, 0.25))

    for k, v in solved_variables.items():
        eq_list2.append(Eq(collector[k][0], float(v)))
        # print(k, '赋值', v)

    variable_collector2 = {
        'number_of_government_workers': (number_of_government_workers, "政府雇员数量", "外生设定"),
        'number_of_total_employment_workers': (number_of_total_employment_workers, "总就业人员数量", "-"),
        'number_of_households': (number_of_households, "居民总量", "外生设定"),
        'ratio_of_dole_to_wage': (ratio_of_dole_to_wage, '救济金与普通工资的比例', "外生设定"),
        'dividends_of_banks': (dividends_of_banks, "银行股利", "-"),
        'net_income_of_households': (net_income_of_households, "居民总收入", "-"),
        'tax_of_households': (tax_of_households, "居民纳税数量", "-"),
        'wages_of_households': (wages_of_households, "就业群体收到工资额", '-'),
        'doles_of_households': (doles_of_households, "失业人员救济金", '-'),
        'deposits_of_households': (deposits_of_households, "居民存款", '-'),
        'deposits_interests_of_households': (deposits_interests_of_households, "居民存款利息", '-'),
        'dividends_of_households': (dividends_of_households, "居民收到股利", '-'),
        'households_income_before_tax_and_dole': (households_income_before_tax_and_dole, "就业居民税前收入", '-'),
        'households_consumption_para_1': (households_consumption_para_1, "居民部门关于收入的消费倾向", '-'),
        'households_consumption_para_2': (households_consumption_para_2, "居民部门关于财富的消费倾向", '外生设定'),
        'real_households_consumption': (real_households_consumption, "居民部门实际消费", '-'),
        'net_wealth_of_households': (net_wealth_of_households, "居民财富", '-'),
        'nominal_households_consumption': (nominal_households_consumption, "居民部门名义消费", '-'),
    }
    variable_collector3 = {**variable_collector2, **collector}

    profits_of_bank_before_tax = Symbol('profits_of_bank_before_tax')
    bank_receive_loans_interests = Symbol('bank_receive_loans_interests')
    interest_rate_of_bonds = Symbol('interest_rate_of_bonds')
    bank_receive_bonds_interests = Symbol('bank_receive_bonds_interests')
    bank_hold_bonds = Symbol('bank_hold_bonds')
    bank_pay_deposits_interests = Symbol('bank_pay_deposits_interests')

    ratio_of_dividends_to_profits_left_in_bank = Symbol('ratio_of_dividends_to_profits_left_in_bank')
    tax_of_bank = Symbol('tax_of_bank')

    interest_rate_of_loans = collector['interest_rate_of_loans'][0]
    loans_of_consumption_firms = collector['loans_of_consumption_firms'][0]
    loan_interests_of_capital_firm = collector['loan_interests_of_capital_firm'][0]
    deposits_interests_of_capital_firm = collector['deposits_interests_of_capital_firm'][0]
    deposits_interests_of_consumption_firm = collector['deposits_interests_of_consumption_firm'][0]

    eq_list3 = []
    loan_interests_of_consumption_firm = Symbol('loan_interests_of_consumption_firm')
    eq_list3.append(
        Eq(loan_interests_of_consumption_firm,
           interest_rate_of_loans * loans_of_consumption_firms * oo1_steady_state_rate)
    )
    eq_list3.append(
        Eq(bank_receive_loans_interests, loan_interests_of_capital_firm + loan_interests_of_consumption_firm)
    )
    eq_list3.append(
        Eq(bank_receive_bonds_interests, interest_rate_of_bonds * bank_hold_bonds * oo1_steady_state_rate)
    )
    eq_list3.append(
        Eq(bank_pay_deposits_interests,
           deposits_interests_of_capital_firm + deposits_interests_of_consumption_firm + deposits_interests_of_households)
    )
    eq_list3.append(
        Eq(profits_of_bank_before_tax,
           bank_receive_loans_interests + bank_receive_bonds_interests - bank_pay_deposits_interests)
    )
    eq_list3.append(
        Eq(tax_of_bank, profits_of_bank_before_tax * tax_rate)
    )
    eq_list3.append(
        Eq(dividends_of_banks, ratio_of_dividends_to_profits_left_in_bank * (profits_of_bank_before_tax - tax_of_bank))
    )

    bank_reserves = Symbol('bank_reserves')
    net_wealth_of_bank = Symbol('net_wealth_of_bank')
    central_bank_hold_bonds = Symbol('central_bank_hold_bonds')

    bonds_government_issue = Symbol('bonds_government_issue')

    loans_of_capital_firms = collector['loans_of_capital_firms'][0]
    deposits_capital_firms_holds = collector['deposits_capital_firms_holds'][0]
    deposits_consumption_firm_holds = collector['deposits_consumption_firm_holds'][0]
    tax_of_capital_firms = collector['tax_of_capital_firms'][0]
    tax_of_consumption_firms = collector['tax_of_consumption_firms'][0]

    eq_list3.append(
        Eq(net_wealth_of_bank,
           loans_of_capital_firms + loans_of_consumption_firms + bank_hold_bonds + bank_reserves -
           deposits_capital_firms_holds - deposits_consumption_firm_holds - deposits_of_households
           )
    )
    eq_list3.append(
        Eq(bank_reserves, central_bank_hold_bonds)
    )
    eq_list3.append(
        Eq(central_bank_hold_bonds, bonds_government_issue - bank_hold_bonds)
    )

    government_pay_bonds_interests = Symbol('government_pay_bonds_interests')
    government_collect_tax = Symbol('government_collect_tax')
    central_bank_profits = Symbol('central_bank_profits')

    eq_list3.append(
        Eq(government_pay_bonds_interests, interest_rate_of_bonds * bonds_government_issue * oo1_steady_state_rate))
    eq_list3.append(Eq(government_collect_tax,
                       tax_of_bank + tax_of_capital_firms + tax_of_consumption_firms + tax_of_households))
    eq_list3.append(Eq(central_bank_profits, interest_rate_of_bonds * central_bank_hold_bonds * oo1_steady_state_rate))
    eq_list3.append(
        Eq(bonds_government_issue * oo2_steady_state_rate,
           households_wage * number_of_government_workers +
           doles_of_households +
           government_pay_bonds_interests -
           government_collect_tax -
           central_bank_profits
           )
    )
    eq_list3.append(Eq(interest_rate_of_bonds, 0.0025))
    eq_list3.append(Eq(ratio_of_dividends_to_profits_left_in_bank, 0.6))

    ratio_of_bank_net_wealth_to_assets = Symbol('ratio_of_bank_net_wealth_to_assets')
    eq_list3.append(Eq(ratio_of_bank_net_wealth_to_assets, 0.08))
    eq_list3.append(
        Eq(net_wealth_of_bank,
           ratio_of_bank_net_wealth_to_assets * (
                   loans_of_capital_firms + loans_of_consumption_firms + bank_hold_bonds + bank_reserves))
    )

    eq_list4 = [*eq_list2, *eq_list3]

    variable_collector4 = {
        'profits_of_bank_before_tax': (profits_of_bank_before_tax, "银行税前利润", "-"),
        'bank_receive_loans_interests': (bank_receive_loans_interests, "银行收到贷款利息", "-"),
        'interest_rate_of_bonds': (interest_rate_of_bonds, "债券利息", "外生设定"),
        'bank_receive_bonds_interests': (bank_receive_bonds_interests, '银行收到债券利息', "-"),
        'bank_hold_bonds': (bank_hold_bonds, "银行持有债券", "-"),
        'bank_pay_deposits_interests': (bank_pay_deposits_interests, "银行支付存款利息", "-"),
        'tax_of_bank': (tax_of_bank, "银行纳税", "-"),
        'bank_reserves': (bank_reserves, "银行准备金", '-'),
        'net_wealth_of_bank': (net_wealth_of_bank, "银行净值", '-'),
        'central_bank_hold_bonds': (central_bank_hold_bonds, "中央银行持有债券", '-'),
        'bonds_government_issue': (bonds_government_issue, "政府发行债券", '-'),
        'government_pay_bonds_interests': (government_pay_bonds_interests, "政府支付债券利息", '-'),
        'government_collect_tax': (government_collect_tax, "政府收税金额", '-'),
        'central_bank_profits': (central_bank_profits, "中央银行利润", '-'),
        'ratio_of_dividends_to_profits_left_in_bank': (ratio_of_dividends_to_profits_left_in_bank, '银行股利比例', '外生设定'),
        'ratio_of_bank_net_wealth_to_assets': (ratio_of_bank_net_wealth_to_assets, "银行净值占资产比例", "-"),
        'loan_interests_of_consumption_firm': (loan_interests_of_consumption_firm, "消费品公司贷款利息", "-"),
    }
    variable_collector4 = {**variable_collector3, **variable_collector4}

    variables4 = [v[0] for k, v in variable_collector4.items()]
    print(f"一共有未知数{len(variables4)}个, 方程{len(eq_list4)}个")
    results = solve(eq_list4, variables4)
    solved_variables = dict(zip([k for k, v in variable_collector4.items()], results[0]))
    set_variables = {k: v for k, v in solved_variables.items() if variable_collector4[k][2] == '外生设定'}
    return solved_variables, set_variables, variable_collector4


def simulate_steady_state():
    dic = {}
    for h in range(50, 250, 5):
        print("模拟steady state rate", h)

        v_set = {'steady_state_rate': h / 10000,
                 'households_wage': 5,
                 'tax_rate': 0.18,
                 }

        solve_v, set_v, v_collector = solve_eq_1_10(v_set)
        solve_v1, set_v1, v_collector1 = solve_eq_11_22(solve_v, v_collector)
        solve_v2, set_v2, v_collector2 = solve_eq_23_30(solve_v1, v_collector1)

        dependent_variables = {k: v for k, v in solve_v2.items() if v_collector2[k][2] != '外生设定'}
        dic[h] = dependent_variables

    df = pd.DataFrame(dic)
    keep_row = [i for i in range(len(df)) if len(set(df.iloc[i, :].to_list())) > 1]
    df_keep = df.iloc[keep_row, :]
    df_keep.index = [v_collector2[s][1] for s in df_keep.index.to_list()]
    # df_keep.T.to_excel('777.xlsx')

    exclude_cols = ['居民财富', '就业居民税前收入']

    df_keep = df_keep.astype(float)
    for i in range(df_keep.shape[1]):
        se = df_keep.iloc[i, :]
        plt.clf()
        a = se.plot()
        plt.title(se.name)
        pth = f"output/sfc_fig/{se.name}.jpg"
        print(pth)
        plt.savefig(pth)


def transfer_to_3d(name, df_keep):
    se = df_keep.T[name]
    df = pd.DataFrame(se)
    r1 =[s.split('_')[0] for s in se.index.to_list()]
    r2 =[s.split('_')[1] for s in se.index.to_list()]
    df['steady_state_rate'] = r1
    df['tax_rate'] = r2
    df = df.astype(float)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 绘制三维表面图
    x = df['steady_state_rate'].values
    y = df['tax_rate'].values
    z = df[name].values
    ax.scatter(x, y, z, c='red', marker='o')

    # 设置坐标轴标签
    ax.set_xlabel('steady_state_rate')
    ax.set_ylabel('tax_rate')
    ax.set_zlabel(name)

    # 显示图形
    return fig


if __name__ == "__main__":
    def simulate_two_para():
        dic = {}
        for h in range(50, 150, 5):
            for tax_rt in range(8, 28, 2):
                print('steady_state:', h, '  tax:',tax_rt)
                v_set = {'steady_state_rate': h / 10000,
                         'households_wage': 5,
                         'tax_rate':tax_rt / 100,
                         }

                solve_v, set_v, v_collector = solve_eq_1_10(v_set)
                solve_v1, set_v1, v_collector1 = solve_eq_11_22(solve_v, v_collector)
                solve_v2, set_v2, v_collector2 = solve_eq_23_30(solve_v1, v_collector1)

                dependent_variables = {k: v for k, v in solve_v2.items() if v_collector2[k][2] != '外生设定'}
                dic[f"{h}_{tax_rt}"] = dependent_variables

                df = pd.DataFrame(dic)
                keep_row = [i for i in range(len(df)) if len(set(df.iloc[i, :].to_list())) > 1]
                df_keep = df.iloc[keep_row, :]
                df_keep.index = [v_collector2[s][1] for s in df_keep.index.to_list()]

        for name in df_keep.index:
            fig = transfer_to_3d(name, df_keep)
            pth = f"output/sfc_fig_2d/{name}.jpg"
            print(pth)
            fig.savefig(pth)

    simulate_two_para()